# -*- coding: utf-8 -*-
# @Time    : 2023/5/8 16:48
# @Author  : Pan
# @Software: PyCharm
# @Project : VisualFramework
# @FileName: DDPMDataset
import os
import cv2
import random
import paddle
import numpy as np
from paddle import io
from datasets import base
from transforms import Compose
from tqdm import tqdm


class DDPMDataset(io.Dataset):
    def __init__(self, config):
        super(DDPMDataset, self).__init__()
        self.mode = config["mode"] if "mode" in config.keys() else "standard"
        self.data_root = config["data_root"] if "data_root" in config.keys() else None
        self.generate_nums = config["generate_nums"] if "generate_nums" in config.keys() else None
        self.recursion_identifier = config["recursion_identifier"] if "recursion_identifier" in config.keys() else None
        self.img_size = config["img_size"] if "img_size" in config.keys() else (224, 224)
        self.trans = Compose(config["transforms"])

        self.data_list = self._make_list(data_root=self.data_root) if self.mode == "standard" else self._make_list(nums=self.generate_nums)

    def __getitem__(self, item):
        data = eval("self._"+self.mode)(item)
        return data

    def __len__(self):
        return len(self.data_list)

    def _standard(self, item):
        data = {"img": self.data_list[item], "path": self.data_list[item]}
        data = self.trans(data)
        return data

    def _predict(self, item):
        return {'img': paddle.randn((3, *self.img_size)), "path": "%08d.png" % item}

    def _make_list(self, data_root=None, nums=None):
        if nums:
            data_list = [i for i in range(nums)]
        else:
            data_list = base.recursion(data_root, self.recursion_identifier)
        return data_list


class TargetDDPMDataset(io.Dataset):
    def __init__(self, config):
        super(TargetDDPMDataset, self).__init__()
        self.mode = config["mode"] if "mode" in config.keys() else "standard"
        self.data_root = config["data_root"] if "data_root" in config.keys() else None
        self.data_list = config["data_list"] if "data_list" in config.keys() else None
        self.generate_nums = config["generate_nums"] if "generate_nums" in config.keys() else None
        self.img_size = config["img_size"] if "img_size" in config.keys() else (224, 224)
        self.trans = Compose(config["transforms"])

        self.data_list = self._make_list(datalist=self.data_list) if self.mode == "standard" else self._make_list(datalist=self.data_list, nums=self.generate_nums)
        self.data_list = self._balance(config["norm"], 256) if "balance" in config.keys() else self.data_list

    def __getitem__(self, item):
        data = eval("self._" + self.mode)(item)
        return data

    def __len__(self):
        return len(self.data_list)

    def _standard(self, item):
        data = {
            "img": os.path.join(self.data_root, self.data_list[item][0]),
            "seg": os.path.join(self.data_root, self.data_list[item][1]),
            "path": self.data_list[item]
        }
        data = self.trans(data)
        return data

    def _predict(self, item):
        data = {'img': paddle.randn((3, *self.img_size)), "seg": os.path.join(self.data_root, self.data_list[item][1]), "path": "%08d.png" % item}
        data = self.trans(data)
        return data

    def _make_list(self, datalist=None, nums=None):
        file = open(datalist, "r")
        list_txt = file.read().replace("\\", "/")
        list_txt = list_txt.rstrip("\n").split("\n")
        file.close()
        list_txt = [t.split(" ") for t in list_txt]
        if nums:
            list_txt = random.choices(list_txt, k=nums)
        return list_txt

    def _balance(self, norm=None, clas=256):
        print("数据平衡开启，正在平衡ing")
        # 得分制的平衡策略
        data_list = []
        # 统计出现的像素个数
        statis = paddle.zeros([clas], dtype=paddle.int64)
        image_statis = paddle.zeros([len(self.data_list), clas], dtype=paddle.int64)
        for idx, item in enumerate(self.data_list):
            target = paddle.to_tensor(cv2.imread(os.path.join(self.data_root, item[1]))).astype(paddle.int64)
            for i in range(clas):
                image_statis[idx, i] = paddle.sum(target == i)
                statis[i] += image_statis[idx, i]

        # 转换成倒数，作为分值，并做归一化
        statis = (1 / statis)
        if norm is not None:
            statis = paddle.log(statis, norm)
        statis = paddle.nan_to_num(statis / statis.min(), 1, 1, 1).astype(paddle.int64)

        # 接下来对所有图片做遍历，即通过各个像素所占的面积作为其出现的次数
        max_repeat = 0
        for idx, item in enumerate(self.data_list):
            repeat = (statis * image_statis[idx]).sum() // image_statis[idx].sum()
            data_list += [item] * repeat.tolist()[0]
            if repeat > max_repeat:
                max_repeat = repeat

        print("根据平衡结果，报告如下：")
        print("raw data list len: %10d" % len(self.data_list))
        print("new data list len: %10d" % len(data_list))
        print("max data list rep: %5d" % max_repeat)
        return data_list
